# Purpose: Defines Pydantic models specific to the agent's internal state and data structures.
# Changes:
# - Added default values/factories for fields that caused instantiation errors (Task, Plan, Step, MemoryRecord, AgentObservation).
# - Made fields truly optional where appropriate (`Optional[...] = None` or `Optional[...] = Field(None, ...)`).

from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional, Literal
import uuid
from datetime import datetime, timezone

# Import dependent schemas
from .mcp.web_automation import PageObservation # Example for AgentObservation data

# --- Task Definition ---
class Task(BaseModel):
    """Represents a task assigned to the agent."""
    id: str = Field(default_factory=lambda: f"task-{uuid.uuid4()}", description="Unique ID for the task.")
    description: str = Field(..., description="The natural language description of the task goal.")
    status: Literal["pending", "running", "completed", "failed", "paused"] = Field("pending", description="Current status of the task.")
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    # Make fields optional or provide defaults if not required at creation
    initial_state: Optional[Dict[str, Any]] = Field(None, description="Optional initial state information provided with the task.")
    result: Optional[Any] = Field(None, description="Final result or outcome of the task.")
    error_message: Optional[str] = Field(None, description="Error message if the task failed.")

# --- Planning ---
class Step(BaseModel):
    """Represents a single action step in an execution plan."""
    id: int = Field(..., description="Sequential step number within the plan (e.g., 1, 2, 3).")
    tool_name: str = Field(..., description="The name of the MCP tool to be called for this step.")
    arguments: Dict[str, Any] = Field(default_factory=dict, description="Arguments required by the MCP tool.")
    thought: Optional[str] = Field(None, description="The LLM's reasoning or thought process behind choosing this step.")
    # Provide defaults for status fields
    status: Literal["pending", "running", "completed", "failed", "skipped"] = Field("pending", description="Execution status of this step.")
    result: Optional[Any] = Field(None, description="The output/result received from the MCP tool after execution.")
    error_message: Optional[str] = Field(None, description="Error message if the step execution failed.")

class Plan(BaseModel):
    """Represents the sequence of steps devised by the Planner to achieve a task."""
    id: str = Field(default_factory=lambda: f"plan-{uuid.uuid4()}", description="Unique ID for the plan.")
    task_id: str = Field(..., description="The ID of the task this plan is for.")
    steps: List[Step] = Field(default_factory=list, description="The ordered list of steps to execute.")
    # Provide defaults for status fields
    status: Literal["draft", "ready", "running", "completed", "failed", "interrupted"] = Field("draft", description="Overall status of the plan execution.")
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    # Make optional fields explicitly Optional
    original_prompt: Optional[str] = Field(None, description="The prompt used to generate this plan (if applicable).")
    raw_llm_response: Optional[str] = Field(None, description="The raw response from the LLM during planning.")

# --- Observation ---
class AgentObservation(BaseModel):
    """A generic wrapper for observations gathered by the agent."""
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    observation_type: str = Field(..., description="Type of observation (e.g., 'web_page', 'file_system', 'api_response').")
    data: Any = Field(..., description="The actual observation data (e.g., PageObservation object, file listing).")
    # Make summary optional
    summary: Optional[str] = Field(None, description="Optional LLM-generated summary of the observation.")

# --- Memory ---
class MemoryRecord(BaseModel):
    """Represents a single record stored in the agent's memory."""
    id: str = Field(default_factory=lambda: f"mem-{uuid.uuid4()}")
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    type: str = Field(..., description="Type of memory record (e.g., 'observation', 'action_result', 'reflection').")
    content: Any = Field(..., description="The actual data stored in the memory record.")
    # Make importance optional
    importance: Optional[float] = Field(None, description="Estimated importance score (e.g., for retrieval).")
    related_task_id: Optional[str] = Field(None, description="ID of the task this memory relates to.")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Other relevant metadata.")